import pickle
import numpy as np
import os.path as osp
import os
from lilab.openlabcluster_postprocess.s1_merge_3_file import get_assert_1_file

def create_srt_file(subtitles, output_file):
    with open(output_file, 'w') as f:
        for i, subtitle in enumerate(subtitles, start=1):
            start_time = subtitle['start_time']
            end_time = subtitle['end_time']
            text = subtitle['text']
            
            f.write(str(i) + '\n')
            f.write(format_time(start_time) + ' --> ' + format_time(end_time) + '\n')
            f.write(text + '\n')
            f.write('\n')

def format_time(time):
    hours = int(time / 3600)
    minutes = int((time % 3600) / 60)
    seconds = int(time % 60)
    milliseconds = int((time - int(time)) * 1000)

    return '{:02d}:{:02d}:{:02d},{:03d}'.format(hours, minutes, seconds, milliseconds)


def create_subtitles(cluster_id_l:np.ndarray, new_cluster_name) -> list:
    assert cluster_id_l.min()==0
    fps = 30
    min_seg = 15
    if len(new_cluster_name)==cluster_id_l.max():
        cluster_id_l -= 1
    srt_record = []
    cluster_id_last, cluster_start, cluster_len = cluster_id_l[0], 0, 1

    for ipos in range(1, len(cluster_id_l)):
        cluster_id = cluster_id_l[ipos]
        if cluster_id == cluster_id_last:
            cluster_len += 1
        else:
            srt_record.append((cluster_id_last, cluster_start, cluster_len))
            cluster_id_last = cluster_id
            cluster_start = ipos
            cluster_len = 1

    srt_record.append((cluster_id_last, cluster_start, cluster_len))
    srt_record = np.array(srt_record)
    srt_record = srt_record[srt_record[:,0]>=0]
    
    srt_record_valid = srt_record[srt_record[:,2]>min_seg]
    start_time_l = (srt_record_valid[:,1]) / fps
    end_time_l = (srt_record_valid[:,1] + srt_record_valid[:,2]) / fps
    text_l = np.array(new_cluster_name)[srt_record_valid[:,0]]

    subtitles = [{
        'start_time':s,
        'end_time':e,
        'text':t
        } for s,e,t in zip(start_time_l, end_time_l, text_l)]
    return subtitles


project_dir = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32_k36/semiseq2seq_iter0/output/far_ns_with_s_recluster_k36/representitive_k36_filt_perc88'
seqpkl = get_assert_1_file(project_dir+'/*_sequences.pkl')
clippredpkl = get_assert_1_file(project_dir+'/*.clippredpkl')

outdir = osp.join(osp.dirname(seqpkl), 'srt')
os.makedirs(outdir, exist_ok=True)
seqdata = pickle.load(open(seqpkl, 'rb'))
clippreddata = pickle.load(open(clippredpkl, 'rb'))


vnakes = clippreddata['df_clipNames']['vnake'].unique()
seq_vnakes = [f'fps30_{v.replace("-","_")}_startFrame0_blackFirst' for v in vnakes]
assert set(seq_vnakes) < set(seqdata.keys())
outfiles = [osp.join(outdir, f'{v}.srt') for v in vnakes]

cluster_names = clippreddata['cluster_names']

for v, outfile in zip(seq_vnakes, outfiles):
    seq = np.array(seqdata[v])
    seq -= seq.min()
    subtitles = create_subtitles(seq, cluster_names)
    create_srt_file(subtitles, outfile)
